package com.lzq.rabbitmqdemo.java_api.publisher_confirms;

import com.lzq.rabbitmqdemo.java_api.constant.ConnectionConfig;
import com.lzq.rabbitmqdemo.java_api.constant.Constants;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmListener;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;

import java.io.IOException;
import java.util.Collections;
import java.util.SortedSet;
import java.util.TreeSet;

public class PublisherConfirms {

    public static void main(String[] args) throws Exception{
        publishMessagesIndividually();
        publishMessagesInBatch();
        handlePublishConfirmsAsynchronously();
    }


    static void publishMessagesIndividually() throws Exception {
        ConnectionFactory factory = ConnectionConfig.getFactory();
        try (Connection connection = factory.newConnection();
             Channel channel = connection.createChannel()) {
            channel.confirmSelect();
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE_NAME1, true, false, true, null);
            long start = System.currentTimeMillis();
            for (int i = 0; i < Constants.MESSAGE_COUNT; i++) {
                String body = "消息" + i;
                channel.basicPublish("", Constants.PUBLISHER_CONFIRMS_QUEUE_NAME1, null, body.getBytes());
                channel.waitForConfirms(5_000);
            }
            long end = System.currentTimeMillis();
            System.out.format("Published %d messages individually in %d ms%n", Constants.MESSAGE_COUNT, end - start);
        }
    }

    static void publishMessagesInBatch() throws Exception {
        ConnectionFactory factory = ConnectionConfig.getFactory();
        try (Connection connection = factory.newConnection();
             Channel channel = connection.createChannel()) {
            channel.confirmSelect();
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE_NAME2, true, false, true, null);
            int batchSize = 100;
            int outstandingMessageCount = 0;
            long start = System.currentTimeMillis();
            for (int i = 0; i < Constants.MESSAGE_COUNT; i++) {
                String body = "消息" + i;
                channel.basicPublish("", Constants.PUBLISHER_CONFIRMS_QUEUE_NAME2, null, body.getBytes());
                outstandingMessageCount++;
                if (batchSize == outstandingMessageCount) {
                    channel.waitForConfirms(5_000);
                    outstandingMessageCount = 0;
                }
            }
            if (outstandingMessageCount > 0) {
                channel.waitForConfirms(5_000);
            }
            long end = System.currentTimeMillis();
            System.out.format("Published %d messages individually in %d ms%n", Constants.MESSAGE_COUNT, end - start);
        }
    }

    static void handlePublishConfirmsAsynchronously() throws Exception {
        ConnectionFactory factory = ConnectionConfig.getFactory();
        try (Connection connection = factory.newConnection();
             Channel channel = connection.createChannel()) {
            channel.confirmSelect();
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE_NAME3, false, false, true, null);

            SortedSet<Long> confirmSet = Collections.synchronizedSortedSet(new TreeSet<>());
            channel.addConfirmListener(new ConfirmListener() {
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                    if (multiple) {
                        confirmSet.headSet(deliveryTag+1).clear();
                    } else {
                        confirmSet.remove(deliveryTag);
                    }
                }

                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {
                    System.err.format("deliveryTag: %d, multiple: %b%n", deliveryTag, multiple);
                    if (multiple) {
                        confirmSet.headSet(deliveryTag+1).clear();
                    } else {
                        confirmSet.remove(deliveryTag);
                    }
                }
            });
            long start = System.currentTimeMillis();
            for (int i = 0; i < Constants.MESSAGE_COUNT; i++) {
                String body = "消息" + i;
                long nextPublishSeqNo = channel.getNextPublishSeqNo();
                channel.basicPublish("", Constants.PUBLISHER_CONFIRMS_QUEUE_NAME3, null, body.getBytes());
               confirmSet.add(nextPublishSeqNo);
            }
            while (!confirmSet.isEmpty()) {
                Thread.sleep(10);
            }
            long end = System.currentTimeMillis();
            System.out.format("Published %d messages individually in %d ms%n", Constants.MESSAGE_COUNT, end - start);
        }
    }

}
